#include "predict_hsblock.hh"

int main(const int argc, const char* argv[])
{
  if (argc < 4)
  {
    print_help();
    return -1;
  }
  boost::shared_ptr<model_tree_t> tree_ptr;
  boost::shared_ptr<sp_stat_mat_t> pairs_ptr;
  boost::shared_ptr<sp_stat_mat_t> Z_ptr;

  // read the tree
  ifstream ifs(argv[1], ios::in);
  tree_ptr = read_btree(ifs);
  ifs.close();

  // read the network
  pairs_ptr = read_pairs(argv[2], true);
  collapsible_network_t G(*(pairs_ptr.get()));

  // read the Z
  Z_ptr = read_latent_mat(argv[3]);

  model_tree_t& Tree = *(tree_ptr.get());
  sp_stat_mat_t& Z = *(Z_ptr.get());

  //  Z.dump();
  predict(Tree, G, Z);

  return 0;
}

void print_help()
{
  cerr << endl;
  cerr << "./predict_hsblock model.tree train.wpairs Z.latent" << endl;
  cerr << endl;
  cerr << " model.tree   = a fitted model of hsblock" << endl;
  cerr << " train.wpairs = a network used for training" << endl;
  cerr << "                u <tab> v <tab> weight <nl>" << endl;
  cerr << " Z.latent     = latent probability matrix file" << endl;
  cerr << "                u <tab> k <tab> probability <nl>" << endl;
  cerr << endl;
}

// Z[u][k] = probability
void predict(const model_tree_t& Tree,       // tree model
    const collapsible_network_t& G, // training network
    const sp_stat_mat_t& Z)        // latent assignment
{
  const double Dtot = G.Dtot;
  cerr << " -- start prediction" << endl;
  int u, v;
  while (cin >> u >> v)
  {
    double score = 0.;
    // \sum_{ab} z_ua z_vb lambda_{ab} * ( d_u * d_v / 2m )
    if (Z.has_row(dim_t(u)) && Z.has_row(dim_t(v)))
    {
#ifdef DEGREE_CUTOFF
      const double d_u = G.degree_sum(dim_t(u));
      const double d_v = G.degree_sum(dim_t(v));
      const double rho = d_u * d_v / Dtot;
#else
      const double rho = 1.;
#endif
      const sp_stat_vec_t& z_u = Z.pull_const_row(dim_t(u));
      const sp_stat_vec_t& z_v = Z.pull_const_row(dim_t(v));

      for (sp_stat_vec_t::iterator_t ki = z_u.begin_nonzero();
          ki != z_u.end_nonzero(); ++ki)
      {
        int k = ki.dim();
        double z_k = ki.value();
        for (sp_stat_vec_t::iterator_t li = z_v.begin_nonzero();
            li != z_v.end_nonzero(); ++li)
        {
          int l = li.dim();
          double z_l = li.value();
          const model_tree_t::node_data_t& F = Tree.get_lca_node_const(k, l);
          score += F.mean() * z_k * z_l;
        }
      }
      score *= rho;
    }
    cout << u << "\t" << v << "\t" << score << endl;
  }
  cerr << " -- done prediction " << endl;
}

boost::shared_ptr<sp_stat_mat_t> read_latent_mat(const char* fname)
{
  boost::shared_ptr<sp_stat_mat_t> ret(new sp_stat_mat_t);
  sp_stat_mat_t& mat = *(ret.get());
  int u, k;
  double w;
  ifstream ifs(fname, ios::in);

  while (ifs >> u >> k >> w)
  {
    mat(dim_t(u), dim_t(k), w);
  }

  ifs.close();

  return ret;
}

